import numpy as np
from scipy.spatial.distance import cdist

def generate_vector_database(num_vectors=1000, dimensions=50):
    """
    生成一个模拟向量数据库
    :param num_vectors: 向量数量
    :param dimensions: 向量维度
    :return: 随机生成的向量数据库
    """
    np.random.seed(42) # 设置随机种子，保证结果可重复
    return np.random.rand(num_vectors, dimensions)

# 创建一个包含 1000 个 50 维向量的数据库
vector_database = generate_vector_database(num_vectors=1000, dimensions=50)

# 2. 模拟查询向量
query_vector = np.random.rand(1, 50) # 生成一个 50 维的查询向量

# 3. 暴力搜索实现
def brute_force_search(query, database, top_k=5):
    """
    暴力搜索实现
    :param query: 查询向量 (1, d)
    :param database: 向量数据库 (N, d)
    :param top_k: 返回最近的前 k 个向量
    :return: 最近的 k 个向量索引及其距离
    """
    # 计算查询向量到数据库中所有向量的欧氏距离
    distances = cdist(query, database, metric='euclidean').flatten()
    # 获取距离最近的前 k 个向量的索引
    top_k_indices = np.argsort(distances)[:top_k]
    # 获取对应的距离
    top_k_distances = distances[top_k_indices]
    return top_k_indices, top_k_distances

# 使用暴力搜索查找最近的 5 个向量
top_k = 5
nearest_indices, nearest_distances = brute_force_search(
    query_vector, vector_database, top_k=top_k)

# 4. 输出结果
print(f"查询向量: {query_vector.flatten()[:5]}...") # 输出查询向量的前5个值
print(f"\n数据库中最近的 {top_k} 个向量索引: {nearest_indices}")
print(f"对应的欧氏距离: {nearest_distances}")